﻿using System;
using System.Collections;
using System.Collections.Generic;
using UnityEngine;
using System.Runtime.InteropServices;
using UnityEngine.Experimental.Rendering;
using UnityEngine.Windows;

public class RTScene
{
    public enum dxrRenderTargetFormat : uint
    {
        Unknown = 0,
        Ru8,
        RGu8,
        RGBAu8,
        Rf16,
        RGf16,
        RGBAf16,
        Rf32,
        RGf32,
        RGBAf32,
    }


    public const string LibName = "dxrRender";

    [DllImport(LibName)]
    static extern void dxrInitSystem();

    [DllImport(LibName)]
    static extern void dxrUnInitSystem();
    [DllImport(LibName)]
    static extern IntPtr dxrCreateMesh();
    [DllImport(LibName)]
    static extern void dxrUpdateCPUMesh(IntPtr mesh, IntPtr vb, IntPtr ib, int vertex_stride, int vertex_count, int vertex_offset,
        int index_stride, int index_count, int index_offset);
    [DllImport(LibName)]
    static extern void dxrUpdateGPUMesh(IntPtr mesh, IntPtr vb, IntPtr ib, int vertex_stride, int vertex_count, int vertex_offset,
        int index_stride, int index_count, int index_offset);
    [DllImport(LibName)]
    static extern IntPtr dxrCreateInstance(IntPtr mesh);
    [DllImport(LibName)]
    static extern void dxrUpdateInstanceTransform(IntPtr inst, Matrix4x4 transform);

    [DllImport(LibName)]
    static extern void dxrCreateCamera();
    [DllImport(LibName)]
    static extern void dxrUpdateCameraTransform(Vector3 pos, Matrix4x4 view, Matrix4x4 proj);
    [DllImport(LibName)]
    static extern IntPtr dxrCreateLight();
    [DllImport(LibName)]
    static extern void dxrUpdateLight(IntPtr light, Vector3 dir);

    [DllImport(LibName)]
    static extern void dxrCreateRenderTarget();
    [DllImport(LibName)]
    static extern void dxrUpdateRenderTarget(IntPtr hostBuffer, int w, int h, dxrRenderTargetFormat format);
    [DllImport(LibName)]
    static extern bool dxrRemoveMesh(IntPtr mesh);
    [DllImport(LibName)]
    static extern bool dxrRemoveInstance(IntPtr inst);
    [DllImport(LibName)]
    static extern bool dxrRemoveCamera();
    [DllImport(LibName)]
    static extern bool dxrRemoveLight(IntPtr light);
    [DllImport(LibName)]
    static extern bool dxrRemoveRenderTarget();
    [DllImport(LibName)]
    static extern bool dxrRender();
    class MeshReference
    {
        public Mesh mesh;
        public int referenceCount;
        public MeshReference(Mesh m)
        {
            mesh = m;
            referenceCount = 1;
        }
    }

    class MeshData
    {
        public IntPtr managedPtr;
    }
    class InstanceData
    {
        public IntPtr managedPtr;
    }
    class LightData
    {
        public IntPtr managedPtr;
    }

    //灯光
    private List<Light> mLights = new List<Light>();
    //mesh数据对应DXR的blas
    //renders数据对应DXR的tlas
    private List<MeshRenderer> mMeshRenders = new List<MeshRenderer>();
    private List<MeshReference> mMeshReference = new List<MeshReference>();

    private Dictionary<Mesh, MeshData> mMeshRecords = new Dictionary<Mesh, MeshData>();
    private Dictionary<MeshRenderer, InstanceData> mInstanceRecords = new Dictionary<MeshRenderer, InstanceData>();
    private Dictionary<Light, LightData> mLightRecords = new Dictionary<Light, LightData>();

    private bool bInitialized = false;
    private static RTScene gInst;
    private RTScene()
    {

    }

    public static RTScene Instance
    {
        get
        {
            if (gInst == null)
                gInst = new RTScene();
            return gInst;
        }
    }

    public void InitSystem()
    {
        dxrInitSystem();
        bInitialized = true;
    }

    public void UnInitSystem()
    {
        dxrUnInitSystem();
        bInitialized = false;
    }

    public void RenderFrame()
    {
        if (bInitialized)
            dxrRender();
    }
    public void ClearSceneData()
    {
        for (int i = mLights.Count - 1; i >= 0; i--)
        {
            RemoveLight(mLights[i]);
        }
        for (int i = mMeshRenders.Count - 1; i >= 0; i--)
        {
            RemoveMeshRender(mMeshRenders[i]);
        }
    }


    public void SetCamera(Camera cam)
    {
        if(cam == null)
        {
            dxrRemoveCamera();
        }
        else
        {
            dxrCreateCamera();
            dxrUpdateCameraTransform(cam.transform.position, cam.worldToCameraMatrix, cam.projectionMatrix);
        }
    }

    public void SetRenderTarget(RenderTexture rt, int w, int h, dxrRenderTargetFormat format)
    {
        if (rt != null)
        {
            dxrCreateRenderTarget();
            dxrUpdateRenderTarget(rt.GetNativeTexturePtr(), w, h, format);
        }
        else
        {
            dxrRemoveRenderTarget();
        }
    }

    public bool AddLight(Light l)
    {
        if (mLights.Contains(l))
            return false;
        mLights.Add(l);
        OnLightAdded(l);
        return true;
    }

    public bool RemoveLight(Light l)
    {
        if (!mLights.Contains(l))
            return false;
        OnLightDelete(l);
        mLights.Remove(l);
        return true;
    }

    public bool UpdateLight(Light l)
    {
        if (!mLights.Contains(l))
            return false;
        LightData ld;
        if (mLightRecords.TryGetValue(l, out ld))
        {
            dxrUpdateLight(ld.managedPtr, l.transform.forward);
            return true;
        }
        return false;
    }

    private bool AddMesh(Mesh mesh)
    {
        foreach(var m in mMeshReference)
        {
            if(m.mesh == mesh)
            {
                m.referenceCount++;
                return false;
            }
        }
        mMeshReference.Add(new MeshReference(mesh));
        OnMeshAdded(mesh);
        return true;
    }

    private bool RemoveMesh(Mesh mesh)
    {
        for (int i = 0; i < mMeshReference.Count; ++i)
        {
            MeshReference t = mMeshReference[i];
            if (t.mesh == mesh)
            {
                if (--t.referenceCount <= 0)
                {
                    OnMeshDelete(mesh);
                    mMeshReference.RemoveAt(i);
                }
                return true;
            }
        }
        return false;
    }


    public bool AddMeshRender(MeshRenderer mr)
    {
        if (mMeshRenders.Contains(mr))
            return false;
        Mesh mesh = mr.GetComponent<MeshFilter>().sharedMesh;
        AddMesh(mesh);
        mMeshRenders.Add(mr);
        OnInstanceAdded(mr, mesh);
        return true;
    }

    public bool RemoveMeshRender(MeshRenderer mr)
    {
        if (!mMeshRenders.Contains(mr))
            return false;
        mMeshRenders.Remove(mr);
        Mesh mesh = mr.GetComponent<MeshFilter>().sharedMesh;
        RemoveMesh(mesh);
        OnInstanceDelete(mr);
        return true;
    }

    public bool UpdateMeshRender(MeshRenderer mr)
    {
        if (!mMeshRenders.Contains(mr))
            return false;
        InstanceData id;
        if (mInstanceRecords.TryGetValue(mr, out id))
        {
            dxrUpdateInstanceTransform(id.managedPtr, mr.transform.localToWorldMatrix);
            return true;
        }
        return false;
    }

    private void OnMeshAdded(Mesh mesh)
    {
        IntPtr nativeMesh = dxrCreateMesh();
        MeshData md = new MeshData();
        md.managedPtr = nativeMesh;
        mMeshRecords.Add(mesh, md);
        //传递mesh buffer数据到native端
        int indexStride = mesh.indexFormat == UnityEngine.Rendering.IndexFormat.UInt16 ? 2 : 4;
        int indexCountMerged = 0;
        int prevIndexEnd = 0;

        //合并子mesh
        int subMeshCount = mesh.subMeshCount;
        for (int smi = 0; smi < subMeshCount; ++smi)
        {
            if (mesh.GetTopology(smi) != MeshTopology.Triangles)
                break;
            int start = (int)mesh.GetIndexStart(smi);
            if (start != prevIndexEnd)
                break;
            int indexCount = (int)mesh.GetIndexCount(smi);
            indexCountMerged += (int)mesh.GetIndexCount(smi);
            prevIndexEnd = start + indexCount;
        }

        dxrUpdateGPUMesh(nativeMesh, mesh.GetNativeVertexBufferPtr(0), mesh.GetNativeIndexBufferPtr(),
            0, mesh.vertexCount, 0, indexStride, indexCountMerged, 0);
    }

    private void OnMeshDelete(Mesh mesh)
    {
        MeshData md;
        if (mMeshRecords.TryGetValue(mesh, out md))
        {
            dxrRemoveMesh(md.managedPtr);
            mMeshRecords.Remove(mesh);
        }
    }

    private void OnInstanceAdded(MeshRenderer mr, Mesh mesh)
    {
        MeshData md;
        if (mMeshRecords.TryGetValue(mesh, out md))
        {
            IntPtr nativeInst = dxrCreateInstance(md.managedPtr);
            InstanceData id = new InstanceData();
            id.managedPtr = nativeInst;
            mInstanceRecords.Add(mr, id);
            dxrUpdateInstanceTransform(nativeInst, mr.transform.localToWorldMatrix);
        }
    }

    private void OnInstanceDelete(MeshRenderer mr)
    {
        InstanceData id;
        if (mInstanceRecords.TryGetValue(mr, out id))
        {
            dxrRemoveInstance(id.managedPtr);
            mInstanceRecords.Remove(mr);
        }
    }

    private void OnLightAdded(Light l)
    {
        IntPtr nativeLight = dxrCreateLight();
        LightData ld = new LightData();
        ld.managedPtr = nativeLight;
        mLightRecords.Add(l, ld);
        dxrUpdateLight(nativeLight, l.transform.forward);
    }

    private void OnLightDelete(Light l)
    {
        LightData ld;
        if (mLightRecords.TryGetValue(l, out ld))
        {
            dxrRemoveLight(ld.managedPtr);
            mLightRecords.Remove(l);
        }
    }
}
